"""
@author: 景云鹏
@email: 310491287@qq.com
@date: 2022/6/26
"""

import random

from algorithms import Algorithm
from benchmarks.function import Function


class Particle:
    def __init__(self, x, swarm):
        self.v = 0
        self.x = x
        self.swarm = swarm
        self.fit = self.swarm.f(x)
        self.best_fit = self.fit
        self.best_x = x.copy()

    def update(self):
        v_ori = self.swarm.w * self.v
        v_local = random.uniform(0, self.swarm.c1) * (self.best_x - self.x)
        v_global = random.uniform(0, self.swarm.c2) * (self.swarm.best_x - self.x)
        self.v = v_ori + v_local + v_global
        self.x += self.v
        self.swarm.f.check_bounds(self.x)
        self.fit = self.swarm.f(self.x)
        if self.fit < self.best_fit:
            self.best_fit = self.fit
            self.best_x = self.x.copy()


class Swarm:
    def __init__(self, f, size, w, c1, c2):
        self.f: Function = f
        self.size = size
        self.w = w
        self.c1 = c1
        self.c2 = c2
        self.particles = []
        self.best_fit = None
        self.best_x = None

    def init_particles(self):
        for _ in range(self.size):
            particle = Particle(self.f.random_x(), self)
            self.particles.append(particle)
        self._update_best_fit()

    def _update_best_fit(self):
        for particle in self.particles:
            if self.best_fit is None or self.best_fit > particle.best_fit:
                self.best_fit = particle.best_fit
                self.best_x = particle.best_x.copy()

    def update_particles(self):
        for particle in self.particles:
            particle.update()
        self._update_best_fit()


class Pso(Algorithm):

    def _init(self):
        self.helper = Swarm(self.f, self.size, 0.8, 2, 2)
        self.helper.init_particles()

    def _update(self):
        self.helper.update_particles()

    def get_point(self):
        x1s = []
        x2s = []
        ys = []
        for particle in self.helper.particles:
            x1s.append(particle.x[0])
            x2s.append(particle.x[1])
            ys.append(particle.fit)
        best = self.helper.best_x[0], self.helper.best_x[1], self.helper.best_fit
        return (x1s, x2s, ys), best

    def get_result(self):
        return self.helper.best_fit, self.helper.best_x
